package structure;

import java.util.LinkedList;
import java.util.logging.Logger;

/**
 * 平衡二叉树
 *
 * @author or2
 * @date 2021/8/9 11:50
 */
public class AVLTree<E extends Comparable<E>> {

    private static final Logger logger = Logger.getLogger(AVLTree.class.getName());

    /**
     * 数结点-数据结构
     */
    private class Node {
        /**
         * 结点的值
         */
        E value;
        /**
         * 结点的左右孩子
         */
        Node lChild, rChild;
        /**
         * 当前树的高度
         */
        int height;

        public Node(E value) {
            this.value = value;
            this.lChild = null;
            this.rChild = null;
            this.height = 1;
        }
    }

    /**
     * 树的根节点
     */
    private Node root;
    /**
     * 当前树的总结点数目
     */
    private int size;

    /**
     * 获取当前结点的高度
     *
     * @param node 检测结点
     * @return 非空时 - 当前结点的高度<br/>为空时空 - 返回高度 0
     */
    private int getHeight(Node node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 当前结点平衡因子是否满足要求
     *
     * @param node 检测结点
     * @return {@code n > 0} - 左子树比右子树高 n<br/>{@code n = 0} - 左子树与右子树一样高<br/>{@code n < 0} - 左子树比右子树矮 n
     */
    private int isBalanced(Node node) {
        int lHeight = getHeight(node.lChild);
        int rHeight = getHeight(node.rChild);
        return lHeight - rHeight;
    }

    /**
     * 向平衡二叉树中添加一个结点
     *
     * @param value 新添加的值
     * @return 自身->链式调用
     */
    public AVLTree<E> add(E value) {
        if (size == 0) {
            root = new Node(value);
            size++;
        } else {
            root = addRecursion(root, value);
        }
        return this;
    }

    /**
     * 递归算法: 检测当前 cNode 的结点添加新元素
     *
     * @param cNode 当前检测的结点
     * @param value 新元素的值
     */
    private Node addRecursion(Node cNode, E value) {
        int compareTo = value.compareTo(cNode.value);
//        不允许重复插入
        if (compareTo == 0) {
            return cNode;
//        若无左右孩子结点, 直接插入返回
        } else if (compareTo < 0 && cNode.lChild == null) {
            cNode.lChild = new Node(value);
            cNode.height = Math.max(getHeight(cNode.lChild), getHeight(cNode.rChild)) + 1;
            size++;
            return cNode;
        } else if (compareTo > 0 && cNode.rChild == null) {
            cNode.rChild = new Node(value);
            cNode.height = Math.max(getHeight(cNode.lChild), getHeight(cNode.rChild)) + 1;
            size++;
            return cNode;
        }
//        按条件递归检测孩子结点
        if (compareTo < 0) {
            cNode.lChild = addRecursion(cNode.lChild, value);
        } else {
            cNode.rChild = addRecursion(cNode.rChild, value);
        }
//        插入后开始更新当前结点高度
        cNode.height = Math.max(getHeight(cNode.lChild), getHeight(cNode.rChild)) + 1;
//        检测当前结点表示的子树是否仍是平衡二叉树
        return balance(cNode);
    }

    /**
     * 计算平衡因子 找出不平衡子树使其平衡
     *
     * @param cNode 当前计算的结点
     * @return 平衡二叉子树的根节点
     */
    private Node balance(Node cNode) {
        int balanceFactor = isBalanced(cNode);
        if (balanceFactor > 1) {
//            LR平衡调整
            if (cNode.lChild != null && isBalanced(cNode.lChild) < 0) {
                cNode.lChild = leftRotate(cNode.lChild);
            }
//            LL平衡调整
            return rightRotate(cNode);

        } else if (balanceFactor < -1) {
//            RL平衡调整
            if (cNode.rChild != null && isBalanced(cNode.rChild) > 0) {
                cNode.rChild = rightRotate(cNode.rChild);
            }
//            RR平衡调整
            return leftRotate(cNode);
        }
        return cNode;
    }

    /**
     * <pre>{@code
     * |————————————————————————————————————————————————————————|
     * |            对结点y进行左旋操作, 返回旋转后的根节点x           |
     * |          y                                 x           |
     * |         / \                               / \          |
     * |       T1   x         向左旋转 (y)         z    y         |
     * |           / \      ———————————————>    / \   / \       |
     * |          T2  z                        T1 T2 T3 T4      |
     * |             / \                                        |
     * |            T3  T4                                      |
     * |————————————————————————————————————————————————————————|
     *  }</pre>
     *
     * @param cNode 最小不符合 平衡二叉树的子树的 根节点
     * @return 旋转后子树的根节点
     */
    private Node leftRotate(Node cNode) {
        Node x = cNode.rChild;
        Node t2 = x.lChild;
//        向左旋转
        x.lChild = cNode;
        cNode.rChild = t2;
//        更新高度
        cNode.height = Math.max(getHeight(cNode.lChild), getHeight(cNode.rChild)) + 1;
        x.height = Math.max(getHeight(x.lChild), getHeight(x.rChild)) + 1;

        return x;
    }

    /**
     * <pre>{@code
     * |————————————————————————————————————————————————————————|
     * |            对结点y进行右旋操作, 返回旋转后的根节点x           |
     * |            y                               x           |
     * |           / \                             / \          |
     * |          x   T4     向右旋转 (y)          z    y         |
     * |         / \       ———————————————>     / \   / \       |
     * |        z   T3                         T1 T2 T3 T4      |
     * |       / \                                              |
     * |      T1  T2                                            |
     * |————————————————————————————————————————————————————————|
     *  }</pre>
     *
     * @param cNode 最小不符合 平衡二叉树的子树的 根节点
     * @return 旋转后子树的根节点
     */
    private Node rightRotate(Node cNode) {
        Node x = cNode.lChild;
        Node t3 = x.rChild;
//        右旋
        x.rChild = cNode;
        cNode.lChild = t3;
//        更新高度
        cNode.height = Math.max(getHeight(cNode.lChild), getHeight(cNode.rChild)) + 1;
        x.height = Math.max(getHeight(x.lChild), getHeight(x.rChild)) + 1;
        return x;
    }

    /**
     * 删除树中值为 {@code value} 的结点
     *
     * @param value 被删除结点的值
     * @return 链式调用, 自身的指针
     */
    public AVLTree<E> remove(E value) {
//        做一下 简单校验
        if (root == null) {
            logger.warning("当前树为空树, 不应执行删除操作!!");
            return this;
        }
//        使用栈模拟函数的递归调用
        LinkedList<Node> removeNodeStack = new LinkedList<>();
//        查找值等于 value 的元素
        Node p = root;
        Node removeNode = null;
        while (p != null) {
            int compare = p.value.compareTo(value);
            if (compare == 0) {
                removeNode = p;
                break;
            } else {
                removeNodeStack.push(p);
                if (compare > 0) {
                    p = p.lChild;
                } else {
                    p = p.rChild;
                }
            }
        }
//        处理 将被删除的结点 (若不存在则打印日志)
        if (removeNode == null) {
            logger.warning("未在树中查找到含有该值的元素!");
            return this;
        }
//        特殊的简便情况: 其没有左右子树 or 没有其中一个子树
        Node parent = null;
        if (removeNode.lChild == null && removeNode.rChild == null) {
            if (removeNode == root) {
                root = null;
                size--;
                return this;
            }
            parent = removeNodeStack.peek();
            assert parent != null;
            //  因其不为根节点, 一定不为空
            if (parent.rChild == removeNode) {
                parent.rChild = null;
            } else {
                parent.lChild = null;
            }
            size--;
            removeBalance(removeNodeStack, this.root);
            return this;
        }
//        接上文: 其既有左子树, 又有右子树 -> 查找该结点的替代结点(左子树的最小结点)
        Node min;
        p = removeNode.rChild;
        LinkedList<Node> minNodeStack = new LinkedList<>();
        while (p != null) {
            if (p.lChild != null) {
                minNodeStack.push(p);
                p = p.lChild;
                continue;
            }
            break;
        }
        min = p;
        //      记得断掉替代节点父亲结点的引用, 否则会循环引用 (先断->循环导致调试时卡死)
        if (!minNodeStack.isEmpty()) {
            minNodeStack.peek().lChild = null;
        }
        //      删除结点的父节点是否为头节点
        if (removeNodeStack.isEmpty()) {
            min.lChild = root.lChild;
            min.rChild = root.rChild;
            root = min;
        } else {
            parent = removeNodeStack.peek();
            min.lChild = removeNode.lChild;
            min.rChild = removeNode.rChild;
            parent.rChild = min;
        }
        size--;

        removeBalance(minNodeStack, removeNodeStack.isEmpty() ? this.root : removeNodeStack.peek());
        removeBalance(removeNodeStack, this.root);
        return this;
    }

    /**
     * 移出结点后， 检测其父节点堆栈， 校验并调整平衡性
     *
     * @param nodes 父节点的堆栈
     * @param root 子树的根节点
     */
    private void removeBalance(LinkedList<Node> nodes, Node root) {
        while (!nodes.isEmpty()) {
            Node pop = nodes.pop();
            Node balance = balance(pop);
            Node peek = nodes.peek();
            if (peek == null) {
                peek = root;
            }
            if (peek.lChild == pop) {
                peek.lChild = balance;
        //    根节点 -> 特殊情况
            } else if (peek.rChild == pop) {
                peek.rChild = balance;
            } else {
                this.root = balance;
            }
        }
    }

    @Override
    public String toString() {
        StringBuilder stringBuilder = new StringBuilder();
        int leafNodesNum = (int) (Math.pow(2, Math.log(size) / Math.log(2)));

        generateString(root, 0, stringBuilder, leafNodesNum);

        return stringBuilder.toString();
    }

    /**
     * 生成方便观看的BST
     *
     * @param root          当前结点做根节点, 不断递归
     * @param depth         当前深度, 每次递归 +1
     * @param stringBuilder 拼接字符串用
     * @param leafNodesNum  防止重复计算, 第一次计算后一直传递
     */
    private void generateString(Node root, int depth, StringBuilder stringBuilder, int leafNodesNum) {
        double section = Math.pow(2, depth) + 1;
        int sectionLength = (int) (leafNodesNum / section);
        String gap = "\t".repeat(sectionLength);

        if (root == null) {
            stringBuilder.append(gap).append("空").append("\n");
            return;
        }

        generateString(root.lChild, depth + 1, stringBuilder, leafNodesNum);
        stringBuilder.append(gap).append(root.value).append("\n");
        generateString(root.rChild, depth + 1, stringBuilder, leafNodesNum);
    }
}
